#ifndef LMM_HPP
#define LMM_HPP

#include <Eigen/Dense>
#include <Eigen/src/Core/Matrix.h>
#include <cmath>
#include <functional>
#include <vector>

#include "GLRK.hpp"
#include "IVPSolver.hpp"

class LMM : public IVPSolver {
public:
  // Constructor
  LMM(int dim, double dt, double t0, double t1, int p)
      : IVPSolver(dim, dt, t0, t1), m_p(p) {
    m_s = p;
    m_u.resize(std::floor((m_t1 - m_t0) / dt) + 1, m_dim);
  }

  // Destructor
  virtual ~LMM() {}

  // Method to perform time integration
  void solve(const Eigen::VectorXd &u0,
             const std::function<Eigen::VectorXd(Eigen::VectorXd, double)> &f) {
    m_u.row(0) = u0; // Set initial condition
    for (int i = 1; i <= m_s; ++i) {
      m_u.row(i) = initialStepSolver(m_u.row(i - 1), m_t0 + (i - 1) * m_dt, f);
    }
    for (int i = m_s + 1; i < m_u.rows(); ++i) {
      int ite = 0;
      m_u.row(i) = m_u.row(i - 1);
      Eigen::VectorXd prevU = m_u.row(i);
      do {
        prevU = m_u.row(i);
        m_u.row(i) = linearMultistepSolver(i, f);
        ite++;
      } while (ite < ITE && (m_u.row(i).transpose() - prevU).norm() > 1e-5);
    }
  }

  // Improved Euler method for initial steps (to be implemented in derived
  // classes if needed)

  // Eigen::VectorXd initialStepSolver(
  //     const Eigen::VectorXd &prevU, const double prevT,
  //     const std::function<Eigen::VectorXd(Eigen::VectorXd, double)> &f) const
  //     {
  //   Eigen::VectorXd result = prevU;
  //   Eigen::VectorXd y1 = f(prevU, prevT);
  //   Eigen::VectorXd y2 = f(prevU + m_dt * y1, prevT + m_dt);
  //   result += m_dt * (y1 + y2) / 2.0;
  //   return result;
  // }
  Eigen::VectorXd initialStepSolver(
      const Eigen::VectorXd &prevU, const double prevT,
      const std::function<Eigen::VectorXd(Eigen::VectorXd, double)> &f) const {
    GLRK glrk(m_dim, m_dt, prevT, prevT + m_dt, 2);
    glrk.solve(prevU, f);
    Eigen::VectorXd result = glrk.getSolution().row(1);
    return result;
  }

  // Unified implementation of linear multistep solver
  Eigen::VectorXd linearMultistepSolver(
      const int n,
      const std::function<Eigen::VectorXd(Eigen::VectorXd, double)> &f) const {
    Eigen::VectorXd result = Eigen::VectorXd::Zero(m_dim);
    const std::vector<double> &alpha =
        getAlpha(); // Get coefficients from derived class
    const std::vector<double> &beta =
        getBeta(); // Get coefficients from derived class

    for (int i = 1; i <= m_s; ++i) {
      result += -alpha[i] * m_u.row(n - i);
      result += m_dt * beta[i] * f(m_u.row(n - i), m_t0 + (n - i) * m_dt);
    }
    result += m_dt * beta[0] * f(m_u.row(n), m_t0 + n * m_dt);
    result /= alpha[0];
    return result;
  }

  // Pure virtual functions to get coefficients
  virtual const std::vector<double> getAlpha() const = 0;
  virtual const std::vector<double> getBeta() const = 0;

  // Getters for time parameters
  double getDt() const { return m_dt; }
  double getT0() const { return m_t0; }
  double getT1() const { return m_t1; }
  double getS() const { return m_s; }

protected:
  // const int m_dim;   // Dimension of the problem
  // const double m_dt; // Time step
  // const double m_t0; // Initial time
  // const double m_t1; // Final time
  const int m_p; // Polynomial order
  int m_s;       // s-step
  // Eigen::MatrixXd m_u; // my solution
};

#endif // TIME_INTEGRATOR_HPP